// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#pragma once

#include <melon/portability.h>
#include <pollux/vector/complex_vector.h>
#include <pollux/vector/flat_vector.h>

namespace kumo::pollux::dwio::common::flatmap {
    namespace detail {
        // Reset vector with the desired size/hasNulls properties
        void reset(VectorPtr &vector, vector_size_t size, bool hasNulls);

        // Reset vector smart pointer if any of the buffers is not single referenced.
        template<typename... T>
        void resetIfNotWritable(VectorPtr &vector, const T &... buffer) {
            MELON_PUSH_WARNING
            MELON_GNU_DISABLE_WARNING("-Wparentheses")
            if ((... || (buffer && buffer->refCount() > 1))) {
                vector.reset();
            }
            MELON_POP_WARNING
        }
    } // namespace detail

    // Struct for keeping track flatmap key stream metrics.
    // Used by keySelectionCallback_ in FlatMapColumnReader
    struct FlatMapKeySelectionStats {
        uint64_t totalKeys = 0;
        uint64_t selectedKeys = 0;
    };

    // Initialize flat vector
    template<typename T>
    void initializeFlatVector(
        VectorPtr &vector,
        memory::MemoryPool &pool,
        const TypePtr &type,
        vector_size_t size,
        bool hasNulls,
        std::vector<BufferPtr> &&stringBuffers = {}) {
        detail::reset(vector, size, hasNulls);
        if (vector) {
            auto &flat_vector = dynamic_cast<FlatVector<T> &>(*vector);
            detail::resetIfNotWritable(vector, flat_vector.nulls(), flat_vector.values());
            if (vector) {
                flat_vector.setStringBuffers(stringBuffers);
            }
        }

        if (!vector) {
            vector = std::make_shared<FlatVector<T> >(
                &pool,
                type,
                hasNulls ? AlignedBuffer::allocate<bool>(size, &pool) : nullptr,
                0 /*length*/,
                AlignedBuffer::allocate<T>(size, &pool),
                std::move(stringBuffers));
            vector->set_null_count(0);
        }
    }

    // Initialize map vector.
    void initializeMapVector(
        VectorPtr &vector,
        const std::shared_ptr<const Type> &type,
        memory::MemoryPool &pool,
        const std::vector<const BaseVector *> &vectors,
        std::optional<vector_size_t> sizeOverride = std::nullopt);

    // Initialize vector with a list of vectors. Make sure the initialized vector
    // has the capacity to hold all data from them.
    void initializeVector(
        VectorPtr &vector,
        const std::shared_ptr<const Type> &type,
        memory::MemoryPool &pool,
        const std::vector<const BaseVector *> &vectors);

    // Copy one value from source vector to target.
    void copyOne(
        const std::shared_ptr<const Type> &type,
        BaseVector &target,
        vector_size_t targetIndex,
        const BaseVector &source,
        vector_size_t sourceIndex);

    // Copy values from source vector to target.
    void copy(
        const std::shared_ptr<const Type> &type,
        BaseVector &target,
        vector_size_t targetIndex,
        const BaseVector &source,
        vector_size_t sourceIndex,
        vector_size_t count);

    // Represent key value based on type.
    template<typename T>
    class KeyValue {
    private:
        T value_;
        size_t h_;

    public:
        explicit KeyValue(T value) : value_{value}, h_{std::hash<T>()(value)} {
        }

        const T &get() const {
            return value_;
        }

        std::size_t hash() const {
            return h_;
        }

        bool operator==(const KeyValue<T> &other) const {
            return value_ == other.value_;
        }
    };

    template<typename T>
    struct KeyValueHash {
        std::size_t operator()(const KeyValue<T> &kv) const {
            return kv.hash();
        }
    };

    template<typename T>
    KeyValue<T> parseKeyValue(std::string_view str) {
        return KeyValue<T>(melon::to<T>(str));
    }

    template<>
    inline KeyValue<StringView> parseKeyValue<StringView>(std::string_view str) {
        return KeyValue<StringView>(StringView(str));
    }

    enum class KeyProjectionMode { ALLOW, REJECT };

    template<typename T>
    struct KeyProjection {
        KeyProjectionMode mode = KeyProjectionMode::ALLOW;
        KeyValue<T> value;
    };

    template<typename T>
    KeyProjection<T> convertDynamic(const melon::Dynamic &v) {
        constexpr char reject_prefix = '!';
        const auto str = v.asString();
        const std::string_view view(str);
        if (!view.empty() && view.front() == reject_prefix) {
            return {
                .mode = KeyProjectionMode::REJECT,
                .value = parseKeyValue<T>(view.substr(1)),
            };
        } else {
            return {
                .mode = KeyProjectionMode::ALLOW,
                .value = parseKeyValue<T>(view),
            };
        }
    }

    template<typename T>
    class KeyPredicate {
    public:
        using Lookup = std::unordered_set<KeyValue<T>, KeyValueHash<T> >;

        KeyPredicate(KeyProjectionMode mode, Lookup keyLookup)
            : keyLookup_{std::move(keyLookup)},
              predicate_{
                  mode == KeyProjectionMode::ALLOW
                      ? &KeyPredicate::allowFilter
                      : &KeyPredicate::rejectFilter
              } {
        }

        bool operator()(const KeyValue<T> &key) const {
            return predicate_(key, keyLookup_);
        }

    private:
        static bool allowFilter(const KeyValue<T> &key, const Lookup &lookup) {
            return lookup.size() == 0 || lookup.count(key) > 0;
        }

        static bool rejectFilter(const KeyValue<T> &key, const Lookup &lookup) {
            return lookup.size() == 0 || lookup.count(key) == 0;
        }

        Lookup keyLookup_;
        std::function<bool(const KeyValue<T> &, const Lookup &)> predicate_;
    };

    template<typename T>
    KeyPredicate<T> prepareKeyPredicate(std::string_view expression) {
        std::vector<KeyProjectionMode> modes;
        std::vector<KeyValue<T> > keys;

        if (!expression.empty()) {
            // JSON parse option?
            auto array = melon::parseJson(expression);
            for (auto v: array) {
                POLLUX_CHECK(!v.isNull(), "map key filter should not be null");
                auto converted = convertDynamic<T>(v);
                modes.push_back(converted.mode);
                keys.push_back(std::move(converted.value));
            }
            VKLOG(1) << "[Flat-Map] key filters count: " << array.size();
        }

        POLLUX_CHECK_EQ(modes.size(), keys.size());
        // You cannot mix allow key and reject key.
        POLLUX_CHECK(
            modes.empty() ||
            std::all_of(modes.begin(), modes.end(), [&modes](const auto& v) {
                return v == modes.front();
                }));

        auto mode = modes.empty() ? KeyProjectionMode::ALLOW : modes.front();

        return KeyPredicate<T>(
            mode, typename KeyPredicate<T>::Lookup(keys.begin(), keys.end()));
    }
} // namespace kumo::pollux::dwio::common::flatmap
